package com.yupi.yuaiagent.chatmemory;

import cn.hutool.json.JSONUtil;
import com.yupi.yuaiagent.domain.dto.RedisMessageDTO;
import org.springframework.ai.chat.memory.ChatMemory;
import org.springframework.ai.chat.messages.AssistantMessage;
import org.springframework.ai.chat.messages.Message;
import org.springframework.ai.chat.messages.SystemMessage;
import org.springframework.ai.chat.messages.UserMessage;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.data.redis.core.RedisTemplate;
import org.springframework.stereotype.Component;

import java.util.ArrayList;
import java.util.List;
import java.util.Map;
import java.util.stream.Collectors;

/**
 * 基于Redis的持久化会话记忆存储
 */
@Component
public class RedisChatMemory implements ChatMemory {

    private RedisTemplate<String, Object> redisTemplate;

    public RedisChatMemory(RedisTemplate<String, Object> redisTemplate) {
        this.redisTemplate = redisTemplate;
    }

    // userId 根据后期自己业务进行替换
    private final String REDIS_CHAT_MEMORY_KEY = "chat_memory:userId:";

    @Override
    public void add(String conversationId, List<Message> messages) {
        for (Message message : messages) {
            long time = System.currentTimeMillis();
            RedisMessageDTO redisMessageDTO = new RedisMessageDTO();
            redisMessageDTO.setContent( message.getText());
            redisMessageDTO.setType(message.getMessageType().toString());
            redisTemplate.opsForHash().put(REDIS_CHAT_MEMORY_KEY + conversationId, String.valueOf(time), redisMessageDTO);
        }
    }

    @Override
    public List<Message> get(String conversationId, int lastN) {
        // 获取该用户的所有键值对
        Map<Object, Object> allMessages = redisTemplate.opsForHash().entries(REDIS_CHAT_MEMORY_KEY + conversationId);
        // 判断该用户该会话id是否有会话记忆
        if (allMessages == null || allMessages.isEmpty()) {
            return List.of();  // 如果没有消息，返回空列表
        }
        List<Message> messages = new ArrayList<>();

        // 按时间戳排序（时间戳是String类型，转换为Long后排序）
        List<Map.Entry<Object, Object>> sortedMessages = allMessages.entrySet().stream()
                .sorted((entry1, entry2) -> {
                    long timestamp1 = Long.parseLong(entry1.getKey().toString());
                    long timestamp2 = Long.parseLong(entry2.getKey().toString());
                    return Long.compare(timestamp2, timestamp1);  // 降序排列
                })
                .collect(Collectors.toList());
        // 获取倒数 N 条消息
        List<Map.Entry<Object, Object>> lastMessages = sortedMessages.stream()
                .limit(lastN)
                .collect(Collectors.toList());


        List<Message> message = lastMessages.stream()
                .map(objectObjectEntry ->{
                    // 反系列化成RedisMessageDTO
                    RedisMessageDTO redisMessageDTO = (RedisMessageDTO) objectObjectEntry.getValue();
                    String type = redisMessageDTO.getType();
                    switch (type){
                        case "USER":
                            return new UserMessage(redisMessageDTO.getContent());
                        case "ASSISTANT":
                            return new AssistantMessage(redisMessageDTO.getContent());
                        case "SYSTEM":
                            return new SystemMessage(redisMessageDTO.getContent());
                        default:
                            throw new IllegalArgumentException("Unknown message type: " + type);
                    }
                })
                .collect(Collectors.toList());
        return message;

    }

    @Override
    public void clear(String conversationId) {
        redisTemplate.delete(REDIS_CHAT_MEMORY_KEY + conversationId);
    }
}
